package com.tlzh.controller;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Map;

@RestController
@RequestMapping("/py")
@CrossOrigin
public class ChatPyController {

    @Resource
    private WebClient webClient1;

    @Value("${spring.ai.ollama.pyUrl}")
    private String pyUrl;

    private final RestTemplate restTemplate = new RestTemplate();

    @GetMapping("/params/query")
    public String query(@RequestParam String prompt) {
        // 构造带有查询参数的URL
        String url = "http://192.168.0.188:8000/generate?input_text=" + URLEncoder.encode(prompt, StandardCharsets.UTF_8);

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 由于请求体为空，我们传递HttpEntity的一个实例，其中请求体为null
        HttpEntity<String> entity = new HttpEntity<>(null, headers);

        // 使用POST方法发送请求并接收响应
        RestTemplate restTemplate = new RestTemplate();
        String response = restTemplate.postForObject(url, entity, String.class);
        return response;
    }

    private static final Logger logger = LoggerFactory.getLogger(ChatPyController.class);

    @PostMapping(value = "/generate_stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> generateStream(@RequestBody Map<String, Object> requestBody) {

        return webClient1.post()
                .uri("/generate_stream")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(requestBody)
                .exchangeToFlux(clientResponse -> {
                    // 打印响应状态码
                    logger.info("Response status code: {}", clientResponse.statusCode());
                    // 打印响应头
                    logger.info("Response headers: {}", clientResponse.headers().asHttpHeaders());
                    return clientResponse.bodyToFlux(String.class)
                            .doOnNext(data -> {
                                // 打印响应数据
                                logger.info("Response data: {}", data);
                                logger.info("Response data length: {}", data.length());

                            })
                            .flatMap(data -> {
                                try {

                                    JSONObject jsonObject = JSON.parseObject(data);
                                    String event1 = jsonObject.getString("event");
                                    String eventData = jsonObject.getString("data");

                                    // 尝试将数据解析为 ServerSentEvent
                                    ServerSentEvent<String> event = ServerSentEvent.<String>builder()
                                            .data(data)
                                            .build();
                                    return Flux.just(event);
                                } catch (Exception e) {
                                    // 解析失败，记录错误日志
                                    logger.error("Failed to parse data as ServerSentEvent: {}", data, e);
                                    return Flux.empty();
                                }
                            });
                });
    }

}
